#! /usr/bin/env python

""" A helper script that starts and stops necessary components of KEDR
framework.
Currently, KernelStrider is supported, support for other tools will be added
in the future.
"""

import os.path
import os
import argparse
import sys
import subprocess


INSTALL_PREFIX = '@CMAKE_INSTALL_PREFIX@'
MODULES_SUBDIR = 'lib/modules'

# The file to contain the list of started components.
running_file = '/run/kedr-running'


class KModComponent(object):
    """A component that is a kernel module with parameters.

    'name' - name of the module.
    'params' - list of names of the parameters, the parameters will be
        passed to the kernel module when it is loaded into the kernel (if
        set).
    """

    def __init__(self, name, params=None):
        self.name = name

        if params:
            self.params = dict([(par, None) for par in params])
        else:
            self.params = {}

    def start(self, module_path, verbose=False):
        """Start the component, i.e. load the module.

        'module_path' - Path to the directory containing the module.
        Returns False on failure, True if successful.
        """
        path = os.path.join(module_path, self.name + ".ko")

        param_spec = []
        for par in self.params:
            if self.params[par]:
                param_spec.append(par + '=' + params[par])

        cmd = ['insmod', path]
        cmd.extend(param_spec)

        if verbose:
            sys.stdout.write('Starting ' + self.name + '\n')
            sys.stdout.write('Command: ' + ' '.join(cmd) + '\n')

        try:
            ret = subprocess.call(cmd)
            if ret != 0:
                sys.stderr.write(' '.join([
                    'Failed to start the component:',
                    'the command returned %d.' % ret]) + '\n')
                return False

        except OSError as e:
            sys.stderr.write('\n'.join([
                'Failed to execute \'%s\':' % (' '.join(cmd)), str(e), '']))
            return False

        return True

    def stop(self, verbose=False):
        """Stop the component, i.e. unload the module.

        Returns False on failure, True if successful.
        """
        cmd = ['rmmod', self.name]

        if verbose:
            sys.stdout.write('Stopping ' + self.name + '\n')
            sys.stdout.write('Command: ' + ' '.join(cmd) + '\n')

        try:
            ret = subprocess.call(cmd)
            if ret != 0:
                sys.stderr.write(' '.join([
                    'Failed to stop the component:',
                    'the command returned %d.' % ret]) + '\n')
                return False

        except OSError as e:
            sys.stderr.write('\n'.join([
                'Failed to execute \'%s\':' % (' '.join(cmd)), str(e), '']))
            return False

        return True


def create_supported():
    """Create the dictionary of the supported components.

    Returns the dict {name: KModComponent} for the supported components.
    """
    comps = [
        KModComponent('kedr_mem_core',
            ['targets', 'sampling_rate', 'umh_dir']),
        KModComponent('kedr_fh_drd_common'),
        KModComponent('kedr_fh_drd_cdev'),
        KModComponent('kedr_fh_drd_net'),
        KModComponent('kedr_simple_trace_recorder', ['nr_data_pages'])]
    return dict([(c.name, c) for c in comps])


def _process_kernel_strider(supported):
    return [supported['kedr_mem_core'],
            supported['kedr_fh_drd_common'],
            supported['kedr_fh_drd_cdev'],
            supported['kedr_fh_drd_net'],
            supported['kedr_simple_trace_recorder']]


def choose_components(tools, params):
    """Create the components needed by the listed tools.

    Returns the list of the components in the order they should be started.
    """
    known_tools = {'KernelStrider': _process_kernel_strider}
    supported = create_supported()
    components = []

    for tool in tools:
        if tool in known_tools:
            components.extend(known_tools[tool](supported))
        else:
            sys.stderr.write('Unknown tool: ' + tool + '\n')
            return None

    for par in params:
        for comp in components:
            if par in comp.params:
                comp.params[par] = params[par]

    return components


def start(components, module_path, verbose=False):
    """Start all components, in order."""

    if os.path.exists(running_file):
        sys.stderr.write(
            'KEDR tools are already running or were not stopped properly.\n')
        if verbose:
            sys.stdout.write('The special file %s exists. ' % running_file)
            sys.stdout.write(
                'If KEDR tools are not running but just were not stopped ')
            sys.stdout.write(
                'properly before reboot, etc., you can remove ')
            sys.stdout.write('%s and try again.\n' % running_file)
            sys.stdout.write(
                '\'lsmod | grep kedr\' will show if the main components ')
            sys.stdout.write('are loaded.\n')
        return False

    failed = False

    # LZO compression routines may be implemented in a kernel module rather
    # than built in. Try to load the module anyway. If it is built in or is
    # already loaded, 'modprobe lzo_compress' will be a no-op.
    cmd = ['modprobe', 'lzo_compress']
    if verbose:
        sys.stdout.write(
            'Trying to load \'lzo_compress\' (may be built-in)\n')
        sys.stdout.write('Command: ' + ' '.join(cmd) + '\n')

    try:
        ret = subprocess.call(cmd)
        if ret != 0:
            sys.stderr.write(' '.join([
                'Failed to load lzo_compress:',
                'the command returned %d.\n' % ret]))
            return False

    except OSError as e:
        sys.stderr.write('\n'.join([
            'Failed to execute \'%s\':' % (' '.join(cmd)), str(e), '']))
        return False


    rollback = []
    for comp in components:
        if comp.start(module_path, verbose=verbose):
            rollback.insert(0, comp)
        else:
            failed = True
            break

    if failed:
        if verbose:
            sys.stdout.write(
                'Failed to start some components, performing rollback.\n')
        for comp in rollback:
            comp.stop(verbose=verbose)
    else:
        try:
            with open(running_file, "w") as rf:
                for comp in rollback:
                    rf.write(comp.name + '\n')

        except IOError as e:
            sys.stderr.write('\n'.join([
                'Failed to write data to %s:' % running_file, str(e), '']))
            for comp in rollback:
                comp.stop(verbose=verbose)
            failed = True

    return not failed


def stop(verbose=False):
    """Stop all components listed in 'running_file', in order.

    'supported' - the list of supported components.
    """
    is_ok = True
    supported = create_supported()

    try:
        with open(running_file, 'r') as rf:
            for comp in rf:
                comp = comp.strip()
                if comp in supported:
                    if not supported[comp].stop(verbose=verbose):
                        is_ok = False
                else:
                    sys.stderr.write('Unknown component: %s\n' % comp)
                    is_ok = False

    except IOError as e:
        sys.stderr.write('KEDR tools are not running.\n')
        if verbose:
            sys.stdout.write(
                'Failed to open %s:\n%s\n' % (running_file, str(e)))
        is_ok = False

    if is_ok:
        if os.path.exists(running_file):
            os.remove(running_file)

    return is_ok


if __name__ == '__main__':
    desc = 'This script starts and stops KEDR tools.'
    usage = '\n'.join([
        '',
        '\t%(prog)s start [--tools=...] [--targets=...] [other options]',
        '\t%(prog)s stop'])

    parser = argparse.ArgumentParser(description=desc, usage=usage)

    grp_action = parser.add_argument_group('actions')
    grp_action.add_argument('action', metavar='start|stop',
                            choices=['start', 'stop'])

    parser.add_argument(
        '--tools', help='comma-separated list of KEDR tools to start')

    parser.add_argument(
        '--targets', help=' '.join([
            'comma-separated list of target kernel modules;',
            'if not specified, all modules loaded after KEDR tools',
            'will be considered targets']),
        default='*')

    parser.add_argument(
        '--prefix', help=' '.join([
            'install prefix, if different from the default one',
            '(%s)' % INSTALL_PREFIX]),
        default=INSTALL_PREFIX)

    parser.add_argument(
        '--nr_data_pages', metavar='NUM',
        help=''.join(['number of pages for the output buffer,',
                      '2 - 65536, must be a power of 2']))

    parser.add_argument(
        '--sampling_rate', metavar='S',
        help=''.join([
            'sampling rate for memory events, ',
            'controls how often memory events from each block of code are ',
            'skipped (not reported). ',
            'The higher the value, the more events are skipped. ',
            '0 or 1 - 31. If it is 0 (default), sampling is disabled, ',
            'all events will be reported. ',
            'May help handle intense event streams at the cost of ',
            'potentially  missed races.']))

    parser.add_argument(
        '--verbose', help='if specified, more messages will be output',
        action='store_true', default=False)

    parser.add_argument(
        '--umh_dir', help=' '.join([
            'directory where the modules will look for the ',
            'usermode helper scrits (used mostly for testing)']))

    args = parser.parse_args()

    # On some older systems, /run/ may be unavailable, /var/tmp/ will be
    # used instead.
    if not os.path.exists('/run/'):
        running_file = '/var/tmp/kedr-running'

    if args.action == 'start':
        if not args.tools:
            args.tools = 'KernelStrider'

        tools = set(args.tools.split(','))

        # [NB] If new parameters are added, do not forget to list them here.
        params = {'targets': args.targets,
                  'nr_data_pages': args.nr_data_pages,
                  'sampling_rate': args.sampling_rate,
                  'umh_dir': args.umh_dir}

        _, _, release, _, _ = os.uname()
        module_path = os.path.join(args.prefix, MODULES_SUBDIR, release,
                                    'misc')

        if args.verbose:
            sys.stdout.write('Module path: %s\n' % module_path)

        components = choose_components(tools, params)
        if not components:
            sys.stderr.write('Found no components to be started.\n')
            sys.exit(1)

        if not start(components, module_path, verbose=args.verbose):
            sys.exit(1)

    elif args.action == 'stop':
        if not stop(verbose=args.verbose):
            sys.exit(1)

    else:
        sys.stderr.write('Unknown action: %s\n' % args.action)
        sys.exit(1)

